# include <gym> for environment creation
import gym
import gym.spaces as spaces
import yaml

import xml.etree.ElementTree as ET
from src.algorithms.hrl.sub_strategies.sub_strategy import SubStrategy, SubstrategiesManager
from src.algorithms.hrl.sub_strategies.search import SearchStrategy
from src.algorithms.hrl.competitive_selfplay import SelfPlayBase,SelfPlayManager

import environments.sub_air_combat_env as my_env



# 初始化飞机数量和ID列表
def _init_planes(config_path='config/Aircraft_config.yaml'):
    planes_count = 0
    red_planes_id_list = []
    blue_planes_id_list = []
    planes_id_list = []
    # 读取仿真配置文件中初始飞机的数量和ID
    tree = ET.parse(config_path)
    root = tree.getroot()
    planes_count = int(root.find('PlanesNum').text)

    for plane in root.findall('Plane'):
        plane_id = plane.get('ID')
        if plane_id / 10010 == 1:  # 红方飞机ID
            red_planes_id_list.append(plane_id)
        else:  # 蓝方飞机ID
            blue_planes_id_list.append(plane_id)
        planes_id_list.append(plane_id)
    return planes_count, planes_id_list, red_planes_id_list, blue_planes_id_list


class DualAircraftEnv(gym.Env):
    def __init__(self,train_type='search',state_space=None, action_space=None,planes_config_path ='config/Aircraft_config.yaml'):
        super(DualAircraftEnv, self).__init__()
        # 定义 状态空间
        self.train_type = train_type
        try:
            if state_space is not None:
                self.observation_space = spaces.Box(state_space[train_type], dtype=float)
            else:
                print("Error: state_space must be provided for DualAircraftEnv initialization.")
        except KeyError:
            print(f"Error: Invalid train_type '{train_type}' for state_space.")
        
        self.action_space = spaces.Box(low=-1, high=1, shape=(5,), dtype=float)  # 角度，速度，高度，火控，雷达

        # 初始化飞机数量和ID列表
        self.planes_count , self.planes_id_list, self.red_planes_id_list, self.blue_planes_id_list = _init_planes(config_path=planes_config_path)
        # 初始化子策略管理器
        self.substrategies_manager = SubstrategiesManager(self.planes_id_list)


    def reset(self):
        # 重置环境状态

        pass

    def step(self, action):
        # 为每个飞机分配子策略和目标

        # 获取当前飞机策略
        for plane_id in self.planes_id_list:
            # 红蓝机判断
            if plane_id / 10010 == 1:
                self.substrategies_manager.update_sub_strategy(self.train_type, plane_id=plane_id)
            else:
                # 根据本机的当前的策略从从对手策略管理器中获取对应的策略--未实现
                self.substrategies_manager.update_sub_strategy(self.train_type, plane_id=plane_id)
        # 获取当前飞机的动作
        for plane_id in self.planes_id_list:
            sub_strategy = self.substrategies_manager.get_sub_strategy(plane_id)
            if sub_strategy is not None:
                angle,speed,high,fire,radar = sub_strategy.GetAction(action)  # 获取当前飞机的动作

        # 执行动作
        

        # 获取当前飞机的状态
        for plane_id in self.planes_id_list:
            sub_strategy = self.substrategies_manager.get_sub_strategy(plane_id)
            if sub_strategy is not None:
                originaldata = my_env.CEGetDecData(plane_id)                # 获取仿真原始数据
                state = sub_strategy.GetState(origindata=originaldata)      # origindata should be passed if needed
            
        

    def render(self, mode='human'):
        # 渲染环境状态
        pass
    

if __name__ == "__main__":
    pass